
import cv2
from PIL.Image import Image
from sahi import AutoDetectionModel
from sahi.predict import get_prediction, get_sliced_prediction, predict

from ultralytics import YOLO
import datetime

# 打开输入视频
input_path = 'E:/EFY/001无人机图片/1'  # 替换为你的输入视频路径
input_video_path = 'E:/EFY/001无人机数据'  # 替换为你的输入视频路径
output_video_path = 'E:/EFY/002无人机数据'  # 替换为你的输出视频路径


detection_model = AutoDetectionModel.from_pretrained(
    model_type='ultralytics',
    # YOLOv8模型的路径
    model_path='weights/visdrone+mydata100+crowdhuman200.pt',
    # YOLOv8模型的路径
    confidence_threshold=0.5,
    # 设备类型。
    # 如果您的计算机配备 NVIDIA GPU，则可以通过将 'device' 标志更改为'cuda:0'来启用 CUDA 加速；否则，将其保留为'cpu'
    device="cuda:0", # or 'cuda:0'
)


# # With an image path
# result = get_sliced_prediction(
#     input_video_path,
#     detection_model,
#     slice_height=1500,
#     slice_width=1500,
#     overlap_height_ratio=0.2,
#     overlap_width_ratio=0.2,
# )
# result.export_visuals(export_dir=output_video_path)

result = predict(
    model_type="ultralytics",
    model_path='weights/visdrone+mydata100+crowdhuman200.pt',
    model_device="0",  # or 'cuda:0'
    model_confidence_threshold=0.5,
    source=input_path,
    slice_height=600,
    slice_width=600,
    overlap_height_ratio=0.2,
    overlap_width_ratio=0.2,
)




